import arrow
import traceback

from mongoengine import Q

from app.db.models.jobs import Job
from app.db.models.tasks import Task, ChangeCheckBasepage
from app.db.models.alerts import Alert
from app.db.models.asset_tasks import AssetTasks

from app.api.v2.jobs import MANUAL_MAP

from app.libs.utils import is_not_none
from app.libs.utility import get_warning
from app.libs.scanner import BatchRemoveTaskSession, BatchStopTaskSession
from app.libs.enums import TaskStatus, TaskTriggerType
from app.celery.handle_result.utils import get_warning_id

from app.handler.v2_balances import BalancesOne

from app.errors import TaskNotFoundError, JobNotFoundError


class LatestTaskSummary:
    result_id_key = 'vulResultId'
    setting_enable_key = 'vulSettings__enable'
    warnings_key = 'level'

    warning_count_pipeline = [
        {'$project': {'result': 1, 'endTime': 1}},
        {'$group': {'_id': None, 'warningCount': {'$sum': {'$size': '$result.warnings'}},
                    'updateTime': {'$max': '$endTime'}}},
        {'$project': {'_id': 0}}
    ]

    today_warning_count_pipeline = [
        {'$project': {'result': 1, 'endTime': 1}},
        {'$group': {'_id': None, 'todayWarningCount': {'$sum': {'$size': '$result.warnings'}}}},
        {'$project': {'_id': 0}}
    ]

    warning_category_pipeline = [
        {'$project': {'result': 1, 'endTime': 1}},
        {'$unwind': '$result.warnings'},
        {'$group': {'_id': f'$result.warnings.{warnings_key}', 'count': {'$sum': 1}}},
        {'$project': {'_id': 0, 'item': '$_id', 'count': 1}}
    ]

    ignored_count_pipeline = [
        {'$project': {'result.warnings.ignoredTime': 1}},
        {'$unwind': '$result.warnings'},
        {'$match': {'result.warnings.ignoredTime': {'$ne': None}}},
        {'$group': {'_id': None, 'warningIgnoredCount': {'$sum': 1}}},
        {'$project': {'_id': 0}}
    ]
    resolved_count_pipeline = [
        {'$project': {'result.warnings.isResolved': 1}},
        {'$unwind': '$result.warnings'},
        {'$match': {'result.warnings.isResolved': True}},
        {'$group': {'_id': None, 'warningResolvedCount': {'$sum': 1}}},
        {'$project': {'_id': 0}}
    ]

    def __init__(self, uid):
        self.uid = uid
        self.resp = {
            'updateTime': "",
            'enables': 0,
            'warningCount': 0,
            'todayWarningCount': 0,
            'warningResolvedCount': 0,
            'warningIgnoredCount': 0,
            self.warnings_key: []
        }
        self.jobs = Job.objects().filter(uid=self.uid).only(self.result_id_key, "enableMonitor")
        self.task_ids = list(filter(is_not_none, self.jobs.distinct(self.result_id_key)))
        self.tasks = Task.objects().filter(pk__in=self.task_ids)

    def cal_enables(self):
        """
        启用资产数 , 产品说改为统计有监测结果的数量
        :return:
        """
        # 统计启用的资产数
        # query = {
        #     'enableMonitor': True,
        #     self.setting_enable_key: True
        # }
        # count = self.jobs.filter(**query).count()

        # 统计有监测结果的数量
        count = self.tasks.count()
        self.resp['enables'] = count
        return count

    def cal_warning_count(self):
        for _ in self.tasks.aggregate(self.warning_count_pipeline):
            self.resp.update(_)
            break

    def cal_resolved_count(self):
        for _ in self.tasks.aggregate(self.resolved_count_pipeline):
            self.resp.update(_)
            break

    def cal_ignored_count(self):
        for _ in self.tasks.aggregate(self.ignored_count_pipeline):
            self.resp.update(_)
            break

    def cal_category(self):
        category = {}
        for _ in self.tasks.aggregate(self.warning_category_pipeline):
            category[_['item']] = _['count']
        self.resp[self.warnings_key] = category

    def cal_today_warning_count(self):
        tasks = self.tasks.filter(endTime__gt=arrow.now().floor('day').datetime)
        for _ in tasks.aggregate(self.today_warning_count_pipeline):
            self.resp.update(_)
            break

    def summary(self):
        fun_names = [a for a in dir(self) if a.startswith('cal_')]
        for fun_name in fun_names:
            self.__getattribute__(fun_name)()
        return self.resp


class VulLatestTaskSummary(LatestTaskSummary):
    pass


class SslLatestTaskSummary(LatestTaskSummary):
    result_id_key = 'sslResultId'
    setting_enable_key = 'sslSettings__enable'
    warnings_key = 'category'

    warning_category_pipeline = [
        {'$project': {'result': 1, 'endTime': 1}},
        {'$unwind': '$result.warnings'},
        {'$group': {'_id': f'$result.warnings.{warnings_key}', 'count': {'$sum': 1}}},
        {'$project': {'_id': 0, 'item': '$_id', 'count': 1}}
    ]


class SecurityEventLatestTaskSummary(SslLatestTaskSummary):
    result_id_key = 'securityEventResultId'
    setting_enable_key = 'securityEventSettings__enable'
    warnings_key = 'category'


class ContentLatestTaskSummary(SslLatestTaskSummary):
    result_id_key = 'contentResultId'
    setting_enable_key = 'contentSettings__enable'
    warnings_key = 'category'


class HttpLatestTaskSummary(LatestTaskSummary):
    result_id_key = 'httpResultId'
    setting_enable_key = 'securityEventSettings__enable'
    warnings_key = 'category'


class LatestTaskDetail:
    result_id_key = 'vulResultId'
    setting_enable_key = 'vulSettings__enable'
    warnings_key = 'level'

    def __init__(self, uid, search_key='', security_status='', trigger_type='', ex_search_key=None, page=1, per_page=10):
        self.page = page
        self.per_page = per_page
        self.ex_search_key = ex_search_key
        if search_key:
            query = Q(uid=uid) & (Q(targetUrl__icontains=search_key) | Q(note__icontains=search_key))
        else:
            query = Q(uid=uid)
        jobs = Job.objects().filter(query).only('vulResultId', 'targetUrl', 'note')

        query = Q(pk__in=list(filter(is_not_none, jobs.distinct(self.result_id_key))))

        if security_status and trigger_type:
            query = query & Q(securityStatus=security_status) & Q(triggerType=trigger_type)
        elif security_status:
            query = query & Q(securityStatus=security_status)
        elif trigger_type:
            query = query & Q(triggerType=trigger_type)

        self.task_query = query

    def detail(self):
        if self.ex_search_key:
            self.task_query = self.get_ex_task_query()
        tasks = Task.objects().filter(self.task_query).order_by('-endTime') \
            .only('jobId', 'taskId', 'name', 'target', 'taskType', 'status', 'taskSettings', 'taskSessionId',
                  'securityStatus', 'triggerType', 'startTime', 'endTime',
                  'result.warningCount', 'result.warningResolvedCount', 'result.warningUnresolvedCount',
                  f'result.warnings.{self.warnings_key}',
                  'result.warnings.isResolved',
                  'result.warnings.ignoredTime',
                  'alertSettings'
                  )
        page = tasks.paginate(self.page, self.per_page)

        items = list(page.items.as_pymongo())

        job_ids = [item['jobId'] for item in items]

        jobs = Job.objects.filter(pk__in=job_ids).only('targetUrl', 'note', 'enableMonitor', 'vulSettings.enable')
        jobs_in_bulk = {}
        for job in jobs:
            jobs_in_bulk[job.pk] = job

        for item in items:
            job = jobs_in_bulk[item['jobId']]
            item['targetUrl'] = job.targetUrl
            item['note'] = job.note
            item['enable'] = job.enableMonitor & job.vulSettings.enable
            # item['level'] = {}

            warnings = item['result'].pop('warnings', [])
            warning_dict = {}
            for warning in warnings:
                # 这里要转换一下,warnings_key, 因为有可能int类型会被转换为浮点类型
                warnings_key = warning.get(self.warnings_key, '')
                if isinstance(warnings_key, float) or isinstance(warnings_key, int):
                    warnings_key = str(int(warnings_key))
                level = warning_dict.setdefault(warnings_key,
                                                {"warningCount": 0, "warningResolvedCount": 0,
                                                 "warningIgnoredCount": 0})
                level['warningCount'] += 1
                if warning.get('isResolved', False):
                    level['warningResolvedCount'] += 1
                if warning.get('ignoredTime'):
                    level['warningIgnoredCount'] += 1
            item['result'][self.warnings_key] = warning_dict

        return {
            'items': items,
            'pageInfo': page.page_info()
        }

    def get_ex_task_query(self):
        return self.task_query & Q(**{f"result__warnings__{self.warnings_key}": self.ex_search_key})


class VulLatestTaskDetail(LatestTaskDetail):
    pass


class SslLatestTaskDetail(LatestTaskDetail):
    result_id_key = 'sslResultId'
    setting_enable_key = 'sslSettings__enable'
    warnings_key = 'level'


class SecurityEventLatestTaskDetail(LatestTaskDetail):
    result_id_key = 'securityEventResultId'
    setting_enable_key = 'securityEventSettings__enable'
    warnings_key = 'category'


class ContentLatestTaskDetail(LatestTaskDetail):
    result_id_key = 'contentResultId'
    setting_enable_key = 'contentSettings__enable'
    warnings_key = 'category'


class BaseEvidence(object):
    common_keys = (
        'recommendation',
        'reference',
        'description',
        'levelName',
        'title',
        'category',
        'impact',
        'foundAt',
        'level',
        'detailText',
        'affects'
    )
    specific_keys = ()

    def __init__(self, task_id=None, warning_id=None, uid=None, ignore_warning=None, addition=None):
        self.addition = addition or {}
        if not ignore_warning:
            query = {'_id': task_id}
            if uid:
                query['uid'] = uid
            task = Task.objects.find(query).first()

            if not task:
                raise TaskNotFoundError()
            job = Job.objects.find({"_id": task.jobId}).first()
            if not job:
                raise JobNotFoundError(str(task.jobId))
            self.task = task
            self.job = job
            self.warning = get_warning(warning_id, task.result.warnings)
        else:
            self.task = None
            self.job = None
            self.warning = ignore_warning
        if self.warning:
            self.results = {k: getattr(self.warning, k, None) for k in self.common_keys}
        else:
            self.results = {k: '' for k in self.common_keys}

    def get(self):

        if self.warning:
            results = {k: getattr(self.warning, k, None) for k in self.specific_keys}
        else:
            results = {k: '' for k in self.specific_keys}

        if (func := getattr(self, 'special_get', None)) and (extra_res := func()):
            self.results.update(extra_res)
        if (func := getattr(self, 'special_get_new', None)) and (extra_res := func()):
            self.results.update(extra_res)
        self.results.update(results)
        if self.job:
            self.results.update({"cydEnabled": self.job.cydEnabled})
        return self.results

    def special_get(self) -> dict:
        pass


class VulEvidence(BaseEvidence):
    specific_keys = ('cnvd', 'cve', 'cnnvd', 'traffics', 'payload')


class SSLEvidence(BaseEvidence):
    def __init__(self, *args, **kwargs):
        super(SSLEvidence, self).__init__(*args, **kwargs)
        cert_info = {}
        if self.task:
            cert_info = self.task.result.addition.get('cert_info', {})
        elif self.addition:
            cert_info = self.addition.get('cert_info', {})
        self.results['certInfo'] = cert_info


class SecurityEventEvidence(BaseEvidence):
    specific_keys = ('affects', 'extraDetail', 'traffics')

    def get(self):
        results = super(SecurityEventEvidence, self).get()
        if (not results.get('foundAt')) and self.task:
            results['foundAt'] = self.task.endTime

        return results

    def special_get(self):
        if func := getattr(self, f"get_{self.warning.category}", None):
            return {'links': func()}

    def special_get_new(self):
        if func := getattr(self, f"get_{self.warning.category}_new", None):
            return {'results': func()}

    def get_cryjack(self) -> list:
        if links := self.warning.detail:
            return links
        return []

    def get_black_links(self):
        return self.warning.detail.get('links', [])

    def get_broken_links(self):
        if links := self.warning.detail:
            return links
        return []

    def get_malscan(self):
        if links := self.warning.detail:
            return links
        return []

    def get_foreign_links(self):
        if links := self.warning.detail:
            return links
        return []

    def get_black_links_new(self):
        return self.warning.detail.get('results', [])


class ContentEvidence(BaseEvidence):
    def get(self):
        results = super(ContentEvidence, self).get()
        if (not results.get('foundAt')) and self.task:
            results['foundAt'] = self.task.endTime

        return results

    def special_get(self) -> dict:
        if not (results := self.warning.addition.get('results', [])):
            return {'detail': getattr(self.warning, 'detail', {}), 'results': []}
        return {'results': results}


class ChangeCheckEvidence(BaseEvidence):

    def get(self):
        results = super(ChangeCheckEvidence, self).get()
        results.update({"detail": self.warning.detail})
        if not results.get("foundAt") and self.task:
            results.update({"foundAt": self.task.endTime})
        return results


class CloudMonitorTasks(object):
    related_collections = (Alert, )

    def __init__(self, task_id_list, ex_query=None, asset_task_map: dict = None):
        if not ex_query:
            ex_query = {}
        self.obj_tasks = Task.objects.filter(pk__in=task_id_list, **ex_query)
        self.task_name_list = []
        self.task_target_list = []
        self.asset_task_map = asset_task_map or {}

        for item in self.obj_tasks:
            self.task_name_list.append(item.name)
            self.task_target_list.append(item.target)

    def delete(self):
        task_oid_list = []
        task_session_id_list = []
        uid_task_type_list = []
        asset_task_list = []
        for obj_task in self.obj_tasks:
            task_oid_list.append(obj_task.pk)
            task_session_id_list.append(obj_task.taskSessionId)
            asset_task_list.append(obj_task.refId)

            if obj_task.triggerType == TaskTriggerType.manual.value \
                    and obj_task.status in [TaskStatus.waiting.value, TaskStatus.active.value]:
                uid_task_type_list.append((obj_task.uid, obj_task.taskType))

        for item in self.related_collections:
            item.objects.filter(taskId__in=task_oid_list).delete()
        BatchRemoveTaskSession(task_session_id_list).batch_remove()
        AssetTasks.objects(pk__in=asset_task_list).delete()
        self.obj_tasks.delete()

        for uid, task_type in uid_task_type_list:
            BalancesOne(uid=uid).inc_manual_task_used(task_type, -1)

    def stop(self):
        task_session_id_list = []
        uid_task_type_list = []
        for obj_task in self.obj_tasks:
            if obj_task.triggerType != TaskTriggerType.manual.value:
                continue
            if obj_task.status not in [TaskStatus.waiting.value, TaskStatus.active.value]:
                continue
            if obj_task.progress >= 99:
                continue
            task_session_id_list.append(obj_task.taskSessionId)
            uid_task_type_list.append((obj_task.uid, obj_task.taskType))
        if task_session_id_list:
            BatchStopTaskSession(task_session_id_list).batch_stop()
            Task.objects.filter(taskSessionId__in=task_session_id_list).update(status=TaskStatus.stopped.value)
            for uid, task_type in uid_task_type_list:
                BalancesOne(uid=uid).inc_manual_task_used(task_type, -1)

    def retry(self):
        for task in self.obj_tasks:
            if task.triggerType != TaskTriggerType.manual.value:
                continue
            if task.status != TaskStatus.stopped.value:
                continue
            balances = BalancesOne(uid=task.uid)
            balances.inc_manual_task_used(task.taskType, 1)
            try:
                schedule = MANUAL_MAP.get(task.taskType)
                if task.taskType == 'hostVul':
                    asset_task = AssetTasks.objects.find_one({'hostVulSettings.taskId': task.taskId, 'uid': task.uid})
                    data = {
                        'hostVulSettings': task.taskSettings.to_dict(),
                        'asset_task_id': str(asset_task.pk),
                        'jobId': task.jobId,
                        'enableMonitor': True,
                        'immediateExec': True,
                        'addition': {'pk': str(task.pk)}
                    }
                else:
                    data = {
                        'taskSettings': task.taskSettings.to_dict(),
                        'jobId': task.jobId,
                        'jobIds': [task.jobId],
                        'enableMonitor': True,
                        'immediateExec': True,
                        'addition': {'pk': str(task.pk)},
                        'name': task.name,
                        'asset_task_id': task.refId,
                        'asset_task': self.asset_task_map.get(task.taskId)
                    }
                schedule(data)
            except Exception:
                traceback.print_exc()
                balances.inc_manual_task_used(task.taskType, -1)


def parse_warning_id(warning, task_obj):
    """根据旧的warning.id 生成新的id"""
    old_id = warning.id
    if "\x01\x02" not in old_id:
        return warning.id
    _, risk_title, _ = warning.id.split("\x01\x02")
    warning_id = get_warning_id(
        task_type=task_obj.taskType,
        job_id=task_obj.jobId,
        url=warning.affects,
        risk_title=risk_title
    )
    return warning_id
